#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/3/19 13:12
@desc: 
"""
import tensorflow as tf


def histogram_equalization_tf(image):
    """
    使用 tensorflow 实现直方图均衡化
    Parameters
    ----------
    image: [height, width, channels]

    Returns
    -------

    """
    shape = tf.shape(image)

    def _equalize_histogram_tf(img):
        values_range = tf.constant([0., 255.], dtype=tf.float32)
        histogram = tf.histogram_fixed_width(
            tf.to_float(img), values_range, 256)
        cdf = tf.cumsum(histogram)
        cdf = tf.to_float(cdf) / tf.to_float(cdf[-1])
        cdf_min = cdf[tf.reduce_min(tf.where(tf.greater(cdf, 0.)))]
        px_map = tf.cast(tf.round((cdf - cdf_min) * 255.), tf.uint8)
        eq_hist = tf.gather_nd(px_map, tf.cast(img, tf.int32))
        return eq_hist

    channels = tf.split(image, 3, axis=2)
    eq_channels = tf.map_fn(_equalize_histogram_tf,
                            tf.convert_to_tensor_or_indexed_slices(channels))
    image = tf.transpose(eq_channels, [1, 2, 0])
    image = tf.reshape(image, shape)
    return image
